import cv2
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.widgets import Slider, Button

plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]


def process_license_plate(image_path):
    image = cv2.imread(image_path)
    if image is None:
        print(f"无法读取图像: {image_path}")
        return

    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 1. 车牌倾斜矫正（霍夫变换）
    def hough_skew_correction(gray_img):
        edges = cv2.Canny(gray_img, 50, 150, apertureSize=3)
        lines = cv2.HoughLines(edges, 1, np.pi / 180, 200)

        if lines is not None:
            angles = []
            for line in lines:
                rho, theta = line[0]
                angles.append(theta * 180 / np.pi)

            median_angle = np.median(angles)
            corrected = rotate_image(gray_img, -median_angle + 90)
            return corrected
        else:
            return gray_img

    def rotate_image(img, angle):
        rows, cols = img.shape
        M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
        rotated = cv2.warpAffine(img, M, (cols, rows), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
        return rotated

    corrected_gray = hough_skew_correction(gray_image)
    corrected_rgb = cv2.cvtColor(corrected_gray, cv2.COLOR_GRAY2RGB)

    # 2. 形态学处理精细定位车牌
    def morphological_processing(gray_img):
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
        opened = cv2.morphologyEx(gray_img, cv2.MORPH_OPEN, kernel, iterations=2)
        closed = cv2.morphologyEx(opened, cv2.MORPH_CLOSE, kernel, iterations=2)
        return closed

    morph_img = morphological_processing(corrected_gray)

    # 3. 垂直投影法去除边框
    def vertical_projection(img):
        hist = cv2.reduce(img, 0, cv2.REDUCE_AVG, dtype=cv2.CV_32F).flatten()
        thresh = np.mean(hist) + 0.3 * np.std(hist)
        left = np.argmax(hist > thresh)
        right = len(hist) - np.argmax(np.flip(hist > thresh))
        return left, right

    def horizontal_projection(img):
        hist = cv2.reduce(img, 1, cv2.REDUCE_AVG, dtype=cv2.CV_32F).flatten()
        thresh = np.mean(hist) + 0.3 * np.std(hist)
        top = np.argmax(hist > thresh)
        bottom = len(hist) - np.argmax(np.flip(hist > thresh))
        return top, bottom

    # 二值化处理用于投影分析
    _, binary = cv2.threshold(morph_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    left, right = vertical_projection(binary)
    top, bottom = horizontal_projection(binary)

    # 裁剪边框
    cropped_rgb = corrected_rgb[top:bottom, left:right]
    cropped_gray = corrected_gray[top:bottom, left:right]

    # 创建图形窗口
    fig, axes = plt.subplots(2, 3, figsize=(18, 10))
    plt.subplots_adjust(bottom=0.35, hspace=0.4)

    # 显示各处理阶段图像
    axes[0, 0].imshow(rgb_image)
    axes[0, 0].set_title('原始图像')
    axes[0, 0].axis('off')

    axes[0, 1].imshow(corrected_rgb, cmap='gray')
    axes[0, 1].set_title('霍夫变换矫正后')
    axes[0, 1].axis('off')

    axes[0, 2].imshow(morph_img, cmap='gray')
    axes[0, 2].set_title('形态学处理后')
    axes[0, 2].axis('off')

    axes[1, 0].imshow(binary, cmap='gray')
    axes[1, 0].set_title('二值化图像')
    axes[1, 0].axis('off')

    axes[1, 1].imshow(cropped_rgb)
    axes[1, 1].set_title('垂直投影裁剪后')
    axes[1, 1].axis('off')

    # HSV阈值处理部分
    hsv_image = cv2.cvtColor(cropped_rgb, cv2.COLOR_RGB2HSV)
    ax_hsv = axes[1, 2]
    hsv_thresholded = ax_hsv.imshow(hsv_image)
    ax_hsv.set_title('HSV阈值处理')
    ax_hsv.axis('off')

    # 初始化HSV阈值
    h_min, s_min, v_min = 0, 0, 0
    h_max, s_max, v_max = 180, 255, 255

    # 创建滑动条
    ax_color = plt.axes([0.25, 0.1, 0.65, 0.03])
    ax_h_min = plt.axes([0.25, 0.25, 0.65, 0.03])
    ax_s_min = plt.axes([0.25, 0.20, 0.65, 0.03])
    ax_v_min = plt.axes([0.25, 0.15, 0.65, 0.03])
    ax_h_max = plt.axes([0.25, 0.10, 0.65, 0.03])
    ax_s_max = plt.axes([0.25, 0.05, 0.65, 0.03])
    ax_v_max = plt.axes([0.25, 0.00, 0.65, 0.03])

    slider_h_min = Slider(ax_h_min, 'H最小值', 0, 180, valinit=h_min)
    slider_s_min = Slider(ax_s_min, 'S最小值', 0, 255, valinit=s_min)
    slider_v_min = Slider(ax_v_min, 'V最小值', 0, 255, valinit=v_min)
    slider_h_max = Slider(ax_h_max, 'H最大值', 0, 180, valinit=h_max)
    slider_s_max = Slider(ax_s_max, 'S最大值', 0, 255, valinit=s_max)
    slider_v_max = Slider(ax_v_max, 'V最大值', 0, 255, valinit=v_max)

    # 更新函数
    def update(val):
        h_min = int(slider_h_min.val)
        s_min = int(slider_s_min.val)
        v_min = int(slider_v_min.val)
        h_max = int(slider_h_max.val)
        s_max = int(slider_s_max.val)
        v_max = int(slider_v_max.val)

        lower = np.array([h_min, s_min, v_min])
        upper = np.array([h_max, s_max, v_max])
        mask = cv2.inRange(hsv_image, lower, upper)
        masked_image = cv2.bitwise_and(cropped_rgb, cropped_rgb, mask=mask)

        hsv_thresholded.set_data(masked_image)
        fig.canvas.draw_idle()

    # 绑定滑动条事件
    for slider in [slider_h_min, slider_s_min, slider_v_min,
                   slider_h_max, slider_s_max, slider_v_max]:
        slider.on_changed(update)

    # 重置按钮
    reset_ax = plt.axes([0.8, 0.35, 0.1, 0.04])
    button = Button(reset_ax, '重置阈值')

    def reset(event):
        slider_h_min.reset()
        slider_s_min.reset()
        slider_v_min.reset()
        slider_h_max.reset()
        slider_s_max.reset()
        slider_v_max.reset()

    button.on_clicked(reset)

    plt.show()


# 替换为实际车牌图像路径
image_path = "3.png"
process_license_plate(image_path)